package com.jianjang.llm.core.filter;

import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.slf4j.MDC;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;
import org.springframework.web.filter.OncePerRequestFilter;

import java.io.IOException;
import java.util.UUID;

@Component
@Order(Ordered.HIGHEST_PRECEDENCE + 10)
public class RequestMdcFilter extends OncePerRequestFilter {

    public static final String MDC_SESSION_ID = "sessionId";
    public static final String HEADER_REQUEST_ID = "X-Request-Id";
    public static final String HEADER_SESSION_ID = "X-Session-Id";

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain)
            throws ServletException, IOException {
        String sessionId = resolveSessionId(request);
        try {
            MDC.put(MDC_SESSION_ID, sessionId);
            // Also reflect back for clients to correlate
            response.setHeader(HEADER_SESSION_ID, sessionId);
            filterChain.doFilter(request, response);
        } finally {
            MDC.remove(MDC_SESSION_ID);
        }
    }

    private String resolveSessionId(HttpServletRequest request) {
        // Priority: explicit headers -> HttpSession -> generated UUID
        String fromHeader = firstNonBlank(request.getHeader(HEADER_SESSION_ID), request.getHeader(HEADER_REQUEST_ID));
        if (isNotBlank(fromHeader)) {
            return fromHeader;
        }
        if (request.getSession(false) != null && isNotBlank(request.getSession(false).getId())) {
            return request.getSession(false).getId();
        }
        return UUID.randomUUID().toString().replace("-", "");
    }

    private String firstNonBlank(String... values) {
        if (values == null) {
            return null;
        }
        for (String v : values) {
            if (isNotBlank(v)) {
                return v;
            }
        }
        return null;
    }

    private boolean isNotBlank(String s) {
        return s != null && !s.trim().isEmpty();
    }
}


